#ifndef TRACKNETV3_TRT_H
#define TRACKNETV3_TRT_H

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <opencv2/opencv.hpp>
#include <vector>
#include <memory>
#include <string>
#include <iostream>

class Logger : public nvinfer1::ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        if (severity <= Severity::kINFO) {
            std::cout << "[TensorRT] " << msg << std::endl;
        }
    }
};

class TrackNetV3_TRT {
public:
    TrackNetV3_TRT();
    ~TrackNetV3_TRT();
    
    bool init(const std::string& engine_path);
    bool process(const std::vector<cv::Mat>& input_frames, cv::Mat& output_heatmap);
    bool destroy();
    
    // 获取模型输入尺寸
    void getInputSize(int& width, int& height, int& channels) const {
        width = input_width_;
        height = input_height_;
        channels = input_channels_;
    }
    
    // 获取模型输出尺寸
    void getOutputSize(int& width, int& height, int& channels) const {
        width = output_width_;
        height = output_height_;
        channels = output_channels_;
    }

    // 从热力图中提取球的位置
    bool getBallPosition(const cv::Mat& heatmap, int& x, int& y, float threshold = 0.3f);
    
private:
    nvinfer1::IRuntime* runtime_;
    nvinfer1::ICudaEngine* engine_;
    nvinfer1::IExecutionContext* context_;
    Logger logger_;
    
    void* input_buffer_;
    void* output_buffer_;

    int input_index_;
    int output_index_;
    
    int max_batch_size_;
    int input_width_;
    int input_height_;
    int input_channels_;
    int output_width_;
    int output_height_;
    int output_channels_;
    
    size_t input_size_;
    size_t output_size_;
    
    cudaStream_t stream_;
    
    void printDims(const nvinfer1::Dims& dims, const std::string& name);
    bool preprocessFrames(const std::vector<cv::Mat>& frames, std::vector<float>& input_data);
    void postprocessOutput(float* output, cv::Mat& heatmap);
};

#endif // TRACKNETV3_TRT_H
